#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
GPU/CPU自适应训练机制演示
展示排列5序列LSTM模型的智能设备识别和自动配置功能
"""

import sys
import os
import torch

# 添加项目根目录到路径
project_root = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, project_root)

def demo_device_adaptive_training():
    """演示GPU/CPU自适应训练机制"""
    print("🖥️ GPU/CPU自适应训练机制演示")
    print("=" * 80)
    
    print(" 系统设备检测:")
    print("-" * 60)
    
    # 1. 检测CUDA可用性
    cuda_available = torch.cuda.is_available()
    print(f"✅ CUDA可用性: {cuda_available}")
    
    if cuda_available:
        # 2. 获取CUDA详细信息
        cuda_version = torch.version.cuda
        gpu_count = torch.cuda.device_count()
        current_device = torch.cuda.current_device()
        gpu_name = torch.cuda.get_device_name(current_device)
        
        print(f"📱 CUDA版本: {cuda_version}")
        print(f"🔢 GPU数量: {gpu_count}")
        print(f"🎯 当前GPU: {current_device}")
        print(f"🏷️ GPU名称: {gpu_name}")
        
        # 3. 显存信息
        gpu_memory = torch.cuda.get_device_properties(current_device).total_memory
        gpu_memory_gb = gpu_memory / (1024**3)
        print(f" GPU显存: {gpu_memory_gb:.2f} GB")
        
        # 4. 当前显存使用情况
        allocated_memory = torch.cuda.memory_allocated() / (1024**3)
        cached_memory = torch.cuda.memory_reserved() / (1024**3)
        print(f"🔋 已分配显存: {allocated_memory:.2f} GB")
        print(f"📦 缓存显存: {cached_memory:.2f} GB")
    else:
        print("⚠️ CUDA不可用，将使用CPU训练")
        print("💡 要启用GPU训练，请安装支持CUDA的PyTorch版本")
    
    print("\\n⚙️ 自适应设备选择逻辑:")
    print("-" * 60)
    
    # 5. 演示自适应设备选择
    device_string = 'cuda' if torch.cuda.is_available() else 'cpu'
    device = torch.device(device_string)
    
    print(f"🎯 自动选择设备: {device}")
    print(f"📝 设备字符串: '{device_string}'")
    print(f"🔧 设备类型: {device.type}")
    
    if device.type == 'cuda':
        print(f"🏷️ GPU索引: {device.index if device.index is not None else 0}")
    
    print("\\n🧪 设备适配测试:")
    print("-" * 60)
    
    try:
        # 6. 创建测试张量验证设备工作
        print("创建测试张量...")
        test_tensor = torch.randn(2, 3, 4).to(device)
        print(f"✅ 张量创建成功，形状: {test_tensor.shape}")
        print(f"📍 张量设备: {test_tensor.device}")
        
        # 7. 测试简单计算
        print("🔢 执行设备计算...")
        result = torch.matmul(test_tensor, test_tensor.transpose(-1, -2))
        print(f"✅ 矩阵乘法成功，结果形状: {result.shape}")
        print(f"📍 结果设备: {result.device}")
        
        # 8. 模拟训练过程的设备使用
        print("🏃 模拟训练步骤...")
        from algorithms.plw_sequence_lstm import PLWSequenceLSTM
        
        # 创建小型模型用于测试
        model = PLWSequenceLSTM(input_dim=5, hidden_dim=32, num_layers=1, dropout=0.1)
        model = model.to(device)
        print(f" 模型已移动到设备: {next(model.parameters()).device}")
        
        # 创建测试输入
        batch_size, seq_len, input_dim = 2, 5, 5
        test_input = torch.randint(0, 10, (batch_size, seq_len, input_dim)).float().to(device)
        print(f"📥 输入张量设备: {test_input.device}")
        
        # 前向传播测试
        with torch.no_grad():
            output = model(test_input)
            print(f"📤 输出张量设备: {output.device}")
            print(f" 输出形状: {output.shape}")
        
        print(" 设备适配测试完全成功！")
        
    except Exception as e:
        print(f" 设备测试失败: {e}")
        return False
    
    print("\\n⚡ 性能优化建议:")
    print("-" * 60)
    
    if cuda_available:
        print("🚀 GPU训练优势:")
        print("  • 并行计算：大幅加速矩阵运算")
        print("  • 大批次：支持更大的batch_size")
        print("  • 快速训练：训练时间大幅缩短")
        print("  • 复杂模型：支持更深层的网络结构")
        
        print("\\n💡 GPU优化建议:")
        print("  • 使用适当的batch_size充分利用GPU")
        print("  • 开启混合精度训练(AMP)节省显存")
        print("  • 使用DataLoader的pin_memory=True")
        print("  • 避免频繁的CPU-GPU数据传输")
    else:
        print("🖥️ CPU训练特点:")
        print("  • 兼容性好：所有环境都支持")
        print("  • 内存灵活：使用系统内存")
        print("  • 稳定性高：不依赖GPU驱动")
        print("  • 成本低：无需专门的GPU硬件")
        
        print("\\n💡 CPU优化建议:")
        print("  • 使用较小的batch_size")
        print("  • 开启多线程数据加载")
        print("  • 考虑使用量化模型")
        print("  • 优化模型结构减少计算量")
    
    print("\\n🔧 实际应用中的自适应配置:")
    print("-" * 60)
    
    print("📋 配置参数自动调整:")
    if cuda_available:
        recommended_config = {
            'batch_size': 32,
            'num_workers': 4,
            'pin_memory': True,
            'epochs': 200,
            'learning_rate': 0.001
        }
    else:
        recommended_config = {
            'batch_size': 16,
            'num_workers': 2,
            'pin_memory': False,
            'epochs': 100,
            'learning_rate': 0.002
        }
    
    for key, value in recommended_config.items():
        print(f"  {key}: {value}")
    
    print("\\n📊 训练时间预估:")
    if cuda_available:
        print("  GPU训练 (7000样本):")
        print("    • 快速训练模式: ~2-5分钟")
        print("    • 完整训练模式: ~10-20分钟")
    else:
        print("  CPU训练 (7000样本):")
        print("    • 快速训练模式: ~10-20分钟")
        print("    • 完整训练模式: ~30-60分钟")
    
    return True

def demo_adaptive_code_examples():
    """演示自适应代码示例"""
    print("\\n💻 自适应代码实现示例:")
    print("=" * 80)
    
    print("🔧 1. 基础设备选择:")
    print("```python")
    print("import torch")
    print("")
    print("# 自动检测并选择最佳设备")
    print("device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')")
    print("print(f'使用设备: {device}')")
    print("```")
    
    print("\\n🔧 2. 模型和数据迁移:")
    print("```python")
    print("# 模型迁移到选定设备")
    print("model = PLWSequenceLSTM().to(device)")
    print("")
    print("# 数据迁移到相同设备")
    print("X_train = torch.FloatTensor(X_train).to(device)")
    print("y_train = torch.LongTensor(y_train).to(device)")
    print("```")
    
    print("\\n🔧 3. 配置参数自适应:")
    print("```python")
    print("CONFIG = {")
    print("    'device': 'cuda' if torch.cuda.is_available() else 'cpu',")
    print("    'batch_size': 32 if torch.cuda.is_available() else 16,")
    print("    'num_workers': 4 if torch.cuda.is_available() else 2,")
    print("    'pin_memory': torch.cuda.is_available()")
    print("}")
    print("```")
    
    print("\\n🔧 4. 运行时设备检查:")
    print("```python")
    print("def train_model():")
    print("    device = torch.device(CONFIG['device'])")
    print("    print(f'📱 使用设备: {device}')")
    print("    ")
    print("    if device.type == 'cuda':")
    print("        print(f'🎯 GPU: {torch.cuda.get_device_name()}')")
    print("        print(f'💾 显存: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f}GB')")
    print("    else:")
    print("        print('🖥️ 使用CPU训练模式')")
    print("```")
    
    print("\\n🔧 5. 错误处理和降级:")
    print("```python")
    print("try:")
    print("    # 尝试使用GPU")
    print("    device = torch.device('cuda')")
    print("    model = model.to(device)")
    print("    test_tensor = torch.randn(1, 1).to(device)")
    print("except Exception as e:")
    print("    print(f'GPU不可用，切换到CPU: {e}')")
    print("    device = torch.device('cpu')")
    print("    model = model.to(device)")
    print("```")

if __name__ == "__main__":
    print("🚀 开始GPU/CPU自适应训练机制演示")
    
    success = demo_device_adaptive_training()
    demo_adaptive_code_examples()
    
    print("\\n" + "=" * 80)
    print("🎯 总结:")
    
    if success:
        print(" GPU/CPU自适应机制工作正常")
        print(" 设备检测和自动配置成功")
        print(" 模型可以在当前环境正常运行")
    else:
        print("⚠️ 部分功能测试失败，但基础机制正常")
    
    print("\\n🔥 关键特性:")
    print("  • 🎯 自动检测：无需手动配置，程序自动识别最佳设备")
    print("  • ⚡ 性能优化：根据设备类型自动调整训练参数")
    print("  • 🛡️ 兼容性强：支持任何PyTorch支持的设备")
    print("  • 🔄 智能降级：GPU不可用时自动使用CPU")
    print("  • 📊 详细反馈：显示设备信息和性能建议")
    
    print("\\n🚀 现在您可以放心使用排列5序列LSTM训练，")
    print("系统会自动为您选择最佳的计算设备！")